#include "utils.h"
#include "libamx.h"

void amx_dct_16x16(int16_t *coefficients)
{
    struct TileConfig cfg_src = {
		1,
		0,
		{0},
		{64,64,64,64,64,64,64,64},
		{0},
		{16,16,16,16,16,16,16,16},
		{0}
    };

    int32_t AMX_ALIGN(64) c[256] = { 0 };
    int32_t AMX_ALIGN(64) d[256] = { 0 };

    uint8_t AMX_ALIGN(64) xh[256] = { 0 };
    uint8_t AMX_ALIGN(64) xl[256] = { 0 };

    uint8_t AMX_ALIGN(64) xh_b[1024] = { 0 };
    uint8_t AMX_ALIGN(64) xl_b[1024] = { 0 };

    // mm512_split_to_two_uint8(coefficients, xh, xl);

    // ToB(xh, xh_b);
    // ToB(xl, xl_b);

    // memcpy(d, HEVC::Sub, sizeof(HEVC::Sub));
    //amx_matrix_mul_zero(&cfg_src, HEVC::Transform_1024, xh_b, c);
    // amx_matrix_mul(&cfg_src, HEVC::Transform_1024, xl_b, d);
    // amx_matrix_mul_hl(&cfg_src, HEVC::Transform_1024, xh_b, xl_b, c, d);
    // mm512_lshift_add_rshift(c, d);

    // DisplayIntergerMatrix(c, 16 * 16, "Result\n", 16);

    return;
}

void dump_tilecfg_reg(struct TileConfig *tcfg, const char* name)
{
    int i;

    printf("\n%s\n", name);
    printf("tcfg[palette_id] = 0x%02x\n", tcfg->palette_id);
    printf("tcfg[start_row ] = 0x%02x\n", tcfg->start_row);
    printf("tcfg[resvd0[14]] = 0x");
    for(i=0; i<14;++i)
    {
        printf("%02x,", tcfg->reserved_0[i]);
    }
    printf("\n");

    printf("tcfg[colsb[8]  ] = 0x");
    for(i=0; i<8;++i)
    {
        printf("%04x,", tcfg->colsb[i]);
    }
    printf("\n");

    printf("tcfg[resvd1[8] ] = 0x");
    for(i=0; i<8;++i)
    {
        printf("%04x,", tcfg->reserved_1[i]);
    }
    printf("\n");

    printf("tcfg[rows[8]   ] = 0x");
    for(i=0; i<8;++i)
    {
        printf("%02x,", tcfg->rows[i]);
    }
    printf("\n");

    printf("tcfg[resvd2[8] ] = 0x");
    for(i=0; i<8;++i)
    {
        printf("%02x,", tcfg->reserved_2[i]);
    }
    printf("\n");
}

uint64_t xgetbv(uint32_t index)
 {
 	uint32_t eax, edx;
 
	asm volatile("xgetbv" : "=a" (eax), "=d" (edx) : "c" (index));
 	return eax + ((uint64_t)edx << 32);
 }

static void cpuid_x86_exec(struct cpuid_t *cpuid, uint32_t eax, uint32_t ecx)
{
    cpuid->ieax = eax;
    cpuid->iecx = ecx;

    asm volatile("pushq %%rbx\n"
        "cpuid\n"
        "movl %%ebx, %1\n"
        "popq %%rbx\n"
        : "=a"(cpuid->eax), "=r"(cpuid->ebx), "=c"(cpuid->ecx), "=d"(cpuid->edx)
        : "a"(cpuid->ieax), "c"(cpuid->iecx)
        : "cc");
}

void do_cpuid_read(uint32_t eax, uint32_t ecx)
{
    struct cpuid_t cpuid;

    /*
     * Test SSE/AVX/FMA
     */
    cpuid_x86_exec(&cpuid, eax, ecx);

    printf("CPUID[%02x:%02x]:eax=[%08x], ebx=[%08x], ecx=[%08x], edx=[%08x]\n",
		    eax, ecx, cpuid.eax, cpuid.ebx, cpuid.ecx, cpuid.edx);
}

void do_cpuid_featuredetect(void)
{
    struct cpuid_t cpuid;

    /*
     * Test SSE/AVX/FMA
     */
    cpuid_x86_exec(&cpuid, 0x1, 0x0);

    if (BIT_TEST(cpuid.edx, 25))
    {
        printf("CPUID : cpu support X86_SSE_ instruction set\n");
    }
    if (BIT_TEST(cpuid.ecx, 28))
    {
        printf("CPUID : cpu support X86_AVX_ instruction set\n");
    }
    if (BIT_TEST(cpuid.ecx, 12))
    {
        printf("CPUID : cpu support X86_FMA_ instruction set\n");
    }

    /*
     * Test AVX512F/AVX512-VNNI
     */
    cpuid_x86_exec(&cpuid, 0x7, 0x0);

    if (BIT_TEST(cpuid.ebx,  5))
    {
        printf("CPUID : cpu support X86_AVX2_ instruction set\n");
    }
    if (BIT_TEST(cpuid.ebx, 16))
    {
        printf("CPUID : cpu support X86_AVX512F_ instruction set\n");
    }
    if (BIT_TEST(cpuid.ecx, 11))
    {
        printf("CPUID : cpu support X86_AVX512_VNNI_ instruction set\n");
    }
    if (BIT_TEST(cpuid.edx, 22))
    {
        printf("CPUID : cpu support X86_AMX-BF16_ instruction set\n");
    }
    if (BIT_TEST(cpuid.edx, 24))
    {
        printf("CPUID : cpu support X86_AMX-TILE_ instruction set\n");
    }
    if (BIT_TEST(cpuid.edx, 25))
    {
        printf("CPUID : cpu support X86_AMX-INT8_ instruction set\n");
    }


    /*
     * Test AVX-VNNI
     */
    cpuid_x86_exec(&cpuid, 0x7, 0x1);

    if (BIT_TEST(cpuid.eax, 4))
    {
        printf("CPUID : cpu support X86_AVX_VNNI_ instruction set\n");
    }
}

#include <unistd.h>
#include <sys/syscall.h>

#define XFEATURE_XTILECFG	17
#define XFEATURE_XTILEDATA	18
#define XFEATURE_MASK_XTILECFG	(1 << XFEATURE_XTILECFG)
#define XFEATURE_MASK_XTILEDATA	(1 << XFEATURE_XTILEDATA)
#define XFEATURE_MASK_XTILE	(XFEATURE_MASK_XTILECFG | XFEATURE_MASK_XTILEDATA)

#define ARCH_GET_XCOMP_PERM	0x1022
#define ARCH_REQ_XCOMP_PERM	0x1023

void req_xtiledata_permission(void)
{
	syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA);
}
